package org.dromara.system.elastic;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.elasticsearch.client.elc.ElasticsearchTemplate;
import org.springframework.data.elasticsearch.core.SearchHit;
import org.springframework.data.elasticsearch.core.SearchHits;
import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
import org.springframework.data.elasticsearch.core.query.SearchTemplateQueryBuilder;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;

public class ElasticsearchUtils {
    private static final Logger logger = LoggerFactory.getLogger(ElasticsearchUtils.class);

    public <T> List<T> multiThreadSearch(ElasticsearchTemplate restTemplate, SearchTemplateQueryBuilder searchQuery, Class<T> clazz, String indexName, int maxThreads, int pageSize) {
        long totalHits = restTemplate.count(searchQuery.build(), IndexCoordinates.of(indexName));
        if(totalHits == 0){
            return new ArrayList<>();
        }
        int totalPages = (int) Math.ceil((double) totalHits / pageSize);
        int numThreads = Math.min(maxThreads, totalPages);

        ExecutorService executorService = Executors.newFixedThreadPool(numThreads);
        CompletionService<List<T>> completionService = new ExecutorCompletionService<>(executorService);

        for (int i = 0; i < totalPages; i++) {
            int page = i;
            completionService.submit(() -> {
                searchQuery.withPageable(PageRequest.of(page, pageSize));

                SearchHits<T> pageSearchHits = restTemplate.search(searchQuery.build(), clazz, IndexCoordinates.of(indexName));
                List<T> pageResults = new ArrayList<>();
                for (SearchHit<T> hit : pageSearchHits.getSearchHits()) {
                    pageResults.add(hit.getContent());
                }
                return pageResults;
            });
        }

        List<T> results = new ArrayList<>();
        for (int i = 0; i < totalPages; i++) {
            try {
                Future<List<T>> future = completionService.take();
                List<T> pageResults = future.get();
                results.addAll(pageResults);
            } catch (InterruptedException | ExecutionException e) {
                logger.error("Error occurred while executing search task", e);
            }
        }

        executorService.shutdown();
        return results;
    }
}

